class FinancialRiskWebSocket {
    constructor(userId) {
        this.userId = userId
        this.socket = null
        this.callbacks = {
            onTaskUpdate: () => {},
            onTaskCompleted: () => {},
            onError: () => {}
        }
    }

    connect() {
        const wsUrl = `ws://${window.location.host}/ws/task/${this.userId}/`
        this.socket = new WebSocket(wsUrl)

        this.socket.onopen = () => {
            console.log('WebSocket connected')
        }

        this.socket.onmessage = (event) => {
            const data = JSON.parse(event.data)
            if (data.type === 'task_update') {
                this.callbacks.onTaskUpdate(data)
                if (data.status === 'COMPLETED') {
                    this.callbacks.onTaskCompleted(data)
                }
            }
        }

        this.socket.onerror = (error) => {
            this.callbacks.onError(error)
        }

        this.socket.onclose = () => {
            console.log('WebSocket disconnected')
        }
    }

    startTask(parameters) {
        if (this.socket && this.socket.readyState === WebSocket.OPEN) {
            this.socket.send(JSON.stringify({
                type: 'start_task',
                parameters: parameters
            }))
        }
    }

    disconnect() {
        if (this.socket) {
            this.socket.close()
        }
    }

    onTaskUpdate(callback) {
        this.callbacks.onTaskUpdate = callback
    }

    onTaskCompleted(callback) {
        this.callbacks.onTaskCompleted = callback
    }

    onError(callback) {
        this.callbacks.onError = callback
    }
}

export default FinancialRiskWebSocket